package com.blogspot.sunserenity;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.SecureRandom;

public class RSAAlgorithm implements AsyncCipher {
	private BigInteger e, d, n;
	private int strength;
	
	public RSAAlgorithm(int strength) throws Exception {
		if (strength < 3) {
			throw new Exception("Wrong RSA bit strength was offered. Chose default 512.");
		}
		this.strength = strength;
		generateKeys();
	}
	
	public void generateKeys() throws Exception{
		BigInteger p, q, phi;
		SecureRandom rand = new SecureRandom();
		do {
			p = BigInteger.probablePrime(strength / 2, rand);
			q = BigInteger.probablePrime(strength / 2, rand);
			phi = (p.subtract(BigInteger.ONE)).multiply(q.subtract(BigInteger.ONE));	// Euler function
			n = p.multiply(q);	// modulus
			int order = 4;		
			do {
				String eNumber = Long.toString((long)Math.pow(2, Math.pow(2, order)) + 1l);
				e = new BigInteger(eNumber); // common practice, one of the Fermat numbers, we starts from 65537 (4th Fermat number)
				order--;
			} while (e.compareTo(phi) != -1 && order >= 0);
			if (order < 0 && e.compareTo(phi) != -1) {
				throw new Exception("Generating public exponent failure. Can't find proper E.");
			} else {
				try {
					d = e.modInverse(phi);
				} catch (Throwable e) {
					e.printStackTrace();
					continue;
				}
			}
		} while (e.gcd(phi).compareTo(BigInteger.ONE) != 0);				
	}
	@Override
	public AsyncKey getPublicKey() {
		return new AsyncKey(e, n);
	}

	/* encrypt with private key */ 
	public byte[] encrypt(byte[] msg) {
		return encrypt (msg, new AsyncKey(d, n));
	}
	
	/* decrypt witn private key*/
	@Override
	public byte[] decrypt(byte[] cipher) {
		return decrypt (cipher, new AsyncKey(d, n));
	}
	@Override
	public byte[] encrypt(byte[] msg, AsyncKey cipherKey) {
		ByteArrayOutputStream baos = new ByteArrayOutputStream();
		final BigInteger nSubstractOne = n.subtract(BigInteger.ONE),
						 e = cipherKey.getExponenta(),
						 n = cipherKey.getModulus();
		BigInteger	m = new BigInteger(msg),
					remainder;
		try {
			while (m.compareTo(BigInteger.ZERO) == 1) {
				remainder = m.subtract(nSubstractOne);
				if (remainder.compareTo(BigInteger.ZERO) == -1) {	// m < n-1
					baos.write(m.modPow(e, n).toByteArray());	//append it to result
				} else { // m >= n-1
					baos.write(nSubstractOne.modPow(e, n).toByteArray()); 	//append it to result
				}
				m = remainder;
			}
		} catch (Throwable exp) {
			exp.printStackTrace();
		}
		return baos.toByteArray();
	}

	public byte[] decrypt(byte[] cipher, AsyncKey cipherKey) {
		final BigInteger d = cipherKey.getExponenta(),
						 n = cipherKey.getModulus();
		BigInteger msg = BigInteger.ZERO;
		ByteArrayInputStream bais = new ByteArrayInputStream(cipher);
		int maxBitCount = n.bitCount();
		byte[] byteBuffer = new byte[maxBitCount];
		try {
			while (bais.available() != 0) {
				int readBytes = bais.read (byteBuffer);
				byte[] biInitBuffer = new byte[readBytes];
					System.arraycopy(byteBuffer, 0, biInitBuffer, 0, biInitBuffer.length); 
						BigInteger bi = new BigInteger (biInitBuffer);
				msg = msg.add(bi.modPow (d, n));
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
		return msg.toByteArray();
	}
}
